import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split
from error_file import NotFittedError


class PolynomialRegression(object):
    """
    多项式回归
    """
    def __init__(self, degree=2, fit_intercept=True,
        copy_X=True,
        n_jobs=None,
        positive=False,
        interaction_only=False,
        include_bias=True,
        order="C"
    ):
        linerregression = LinearRegression(fit_intercept=fit_intercept, copy_X=copy_X, n_jobs=n_jobs, positive=positive)
        polynomialfeatures = PolynomialFeatures(degree=degree, interaction_only=interaction_only,
                                                include_bias=include_bias, order=order)
        self.pipeline = make_pipeline(polynomialfeatures, linerregression)  # 构建自动化流程
        self._is_fitted = False

    def init_data(self, x_train, y_train):
        """
        param: x_train
        param: y_train
        return:
        """
        (m, n) = np.shape(x_train)
        tem_m = np.shape(y_train)
        try:
            if m != tem_m[0]:
                raise ValueError
            self.x_train = x_train
            self.y_train = y_train
        except ValueError:
            print("数据出错，x,y样本不一致")
            return

    def init_set(self):
        pass

    def fit(self, x_train, y_train):
        self.pipeline.fit(x_train, y_train)  # 训练模型
        self._is_fitted = True

    def predict(self, x_test):
        if self._is_fitted:
            self.y_pred = self.pipeline.predict(x_test)  # 在测试集上预测结果并保存在y_pred变量中
            return self.y_pred
        else:
            raise NotFittedError("fit before predict!")

    def get_mse(self, y_test):
        """ 均方误差:MSE """
        return mean_squared_error(y_test, self.y_pred)

    def get_rmse(self, y_test):
        """ 均方根误差:RMSE """
        return np.sqrt(mean_squared_error(y_test, self.y_pred))

    def get_mae(self, y_test):
        """ 平均绝对误差:MAE """
        return mean_absolute_error(y_test, self.y_pred)

    def get_score(self, x_test, y_test):
        """R^2"""
        return self.pipeline.score(x_test, y_test)


if __name__ == '__main__':
    new_pumpkins = pd.read_csv("../test_file/new_pumpkins.csv")  # 利用pandas库打开csv数据
    new_pumpkins.info()
    X = pd.get_dummies(new_pumpkins['Variety']) \
        .join(new_pumpkins['Month']) \
        .join(pd.get_dummies(new_pumpkins['City'])) \
        .join(pd.get_dummies(new_pumpkins['Package']))
    Y = new_pumpkins['Price']
    X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=0)
    polynomial = PolynomialRegression()
    polynomial.fit(X_train, Y_train)
    #
    y_prdedict = polynomial.predict(X_test)
    rmse = polynomial.get_rmse(Y_test)
    score = polynomial.get_score(X_test, Y_test)
    print(f'RMSE指标: {rmse:3.3} ({rmse / np.mean(y_prdedict) * 100:3.3}%)')
    print('相关系数: ', score)





